import argparse
import numpy as np

from collections import defaultdict

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, default="/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/fast5s.CG.call_mods.21mer.tsv", required=False,
                        help="input data")#/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/fast5s.CG.call_mods.21mer.tsv
    parser.add_argument('--bed', type=str, default="/home/xiaoyifu/data/HG002/Bisulfite/non_singleton.bisulfite.bed", required=False,
                        help="bed")
    parser.add_argument('--output', type=str, default="/home/xiaoyifu/data/HG002/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/consensus_non_singleton.tsv", required=False,
                        help="output data")
    return parser.parse_args()
def filter_elements(nums, top,lower):
    filtered = set() 
    
    for num in nums:
        if num <= top and num>=lower:
            filtered.add(num) 
            
    return filtered
def read_file(input_file,input_bed,output_file,prob_cf=0):
    bedsites=defaultdict(set)
    #tsvsites=defaultdict(set)
    predsites={}
    msites={}
    wf = open(output_file, 'w')
    with open(input_file, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            #if words[0] not in sites.keys():
            #    sites[words[0]]=set()
            #tsvsites[(words[0],words[2])].add(int(words[1]))
            test_value=1 if (float(words[7])-float(words[6]))>prob_cf else 0
            key=(words[0],words[2],int(words[1]))
            predsites[key]=test_value
    with open(input_bed, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            #if words[0] not in sites.keys():
            #    sites[words[0]]=set()
            bedsites[(words[0],words[5])].add(int(words[1]))
            test_value= 1 if int(words[10])>50 else 0
            key=(words[0],words[5],int(words[1]))
            msites[key]=test_value
    for chr,strand in bedsites.keys():
        #print(len(sites[chr]))
        for chr_site in bedsites[(chr,strand)]:        
            non_singletons=filter_elements(bedsites[(chr,strand)],chr_site+10,chr_site-10)
            methyls=set()
            pred_methyls=set()
            for non_singleton in non_singletons:
                key=(chr,strand,non_singleton)
                if key not in predsites.keys():
                    continue
                methyls.add(msites[key])
                pred_methyls.add(predsites[key])
            consensus=1 if len(methyls)==1 else 0
            pred_consensus=1 if len(pred_methyls)==1 else 0
            wf.write(chr+'\t'+strand+'\t'+str(chr_site)+'\t'+str(consensus)+'\t'+str(pred_consensus)+'\n')      

if __name__ == '__main__':
    args=parse_args()
    #read_file(args.bed,args.output+'consensus_non_singleton_bed.tsv')
    read_file(args.input,args.bed,args.output)